# coding=utf-8
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, List, Tuple
import torch
import torch.distributed
from torch import nn
from transformers.activations import ACT2FN
from atb_llm.utils.layers.linear import FastLinear
from atb_llm.utils.layers import (
    TensorParallelRowLinear,
    TensorParallelColumnLinear,
    PositionRotaryEmbedding,
    TensorEmbedding,
    load_column_multi,
    paged_attn,
    flash_attn,
    reshape_and_cache
)
from atb_llm.utils.moe_utils import assign
from atb_llm.utils.quantize.pack_type import PackType, calc_linear_pack_type
from atb_llm.utils.log import logger
from atb_llm.utils.log.error_code import ErrorCode
from atb_llm.utils.weights import ProcessGroupType


class DeepseekRMSNorm(nn.Module):
    def __init__(self, prefix, weights, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()

        weight = weights.get_tensor(f"{prefix}.weight")
        self.weight = nn.Parameter(weight)
        self.variance_epsilon = eps

    def forward(self, hidden_states, residual=None):
        if hidden_states.shape[-1] > 8192:
            if residual is not None:
                hidden_states += residual
            residual = hidden_states

        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(
            variance + self.variance_epsilon
        )

        # convert into half-precision if necessary
        if self.weight.dtype in [torch.float16, torch.bfloat16]:
            hidden_states = hidden_states.to(self.weight.dtype)

        return self.weight * hidden_states, residual


class DeepseekMLP(nn.Module):
    def __init__(self, prefix, config, weights, intermediate_size=None):
        super().__init__()
        act = config.hidden_act
        approximate = "tanh" if act in ["gelu_fast", "gelu_pytorch_tanh"] else "none"
        self.act = (
            ACT2FN[act]
            if "gelu" not in act
            else lambda x: torch.nn.functional.gelu(x, approximate=approximate)
        )
        linear_names = [f'{prefix}.up_proj', f'{prefix}.gate_proj']
        pack_name = f'{prefix}.gate_up_proj'
        layer_prefix = '.'.join(prefix.split('.')[:-1])
        norm_name = f'{layer_prefix}.post_attention_layernorm'
        self.pack_type = calc_linear_pack_type(weights, linear_names, norm_name, pack_name)

        if self.pack_type in [
            PackType.ALL_FP, PackType.ALL_W8A8, PackType.ALL_W8A8_ANTI, PackType.ALL_W4A16,
            PackType.ALL_W4A16_ANTI, PackType.ALL_W8A16, PackType.ALL_W8A16_ANTI,
            PackType.MIX_W8A8_DYNAMIC, PackType.MIX_W8A8_DYNAMIC_ANTI,
            PackType.ALL_W8A8_DYNAMIC, PackType.ALL_W8A8_DYNAMIC_ANTI
        ]:
            self.gate_up_proj = load_column_multi(
                config,
                prefixes=[f"{prefix}.gate_proj", f"{prefix}.up_proj"],
                weights=weights,
                head_size=1,
            )
        elif self.pack_type in [PackType.ALL_W8A8SC, PackType.ALL_W8A8SC_ANTI]:
            self.gate_up_proj = TensorParallelColumnLinear.load(
                config,
                prefix=f"{prefix}.gate_up_proj",
                weights=weights,
                bias=False,
            )
        else:
            self.gate_proj = TensorParallelColumnLinear.load(
                config,
                prefix=f"{prefix}.gate_proj",
                weights=weights,
                bias=False,
            )
            self.up_proj = TensorParallelColumnLinear.load(
                config,
                prefix=f"{prefix}.up_proj",
                weights=weights,
                bias=False,
            )
        self.down_proj = TensorParallelRowLinear.load(
            config,
            prefix=f"{prefix}.down_proj",
            weights=weights,
            bias=False,
        )
        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
        self.intermediate_size = (
                (config.intermediate_size + weights.process_group.size() - 1) // weights.process_group.size()
        )

    def forward(self, hidden_states):
        gate_up_states = self.gate_up_proj(hidden_states)
        gate_up_states = gate_up_states.view(-1, 2, self.intermediate_size)
        return self.down_proj(self.act(gate_up_states[:, 0]) * gate_up_states[:, 1])


class FlashDeepseekAttention(torch.nn.Module):
    class ForwardInputArgs:
        def __init__(self,
                     hidden_states: torch.tensor,
                     cos: torch.tensor,
                     sin: torch.tensor,
                     cu_seqlen_prefill: torch.tensor,
                     kv_cache: Tuple[torch.tensor, torch.tensor],
                     block_tables: torch.tensor,
                     slots: torch.tensor,
                     input_lengths: torch.tensor,
                     max_s: torch.tensor):
            self.hidden_states = hidden_states
            self.cos = cos
            self.sin = sin
            self.cu_seqlen_prefill = cu_seqlen_prefill
            self.kv_cache = kv_cache
            self.block_tables = block_tables
            self.slots = slots
            self.input_lengths = input_lengths
            self.max_s = max_s

    def __init__(
            self,
            prefix: str,
            config,
            weights,
    ):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_size = self.hidden_size // self.num_heads

        self.rotary_emb = PositionRotaryEmbedding.static(dim=self.head_size, base=10000.0, device="cpu").to(
            weights.device)

        self.softmax_scale = self.head_size ** -0.5
        if (config.num_attention_heads != config.num_key_value_heads and
            self.num_heads % weights.process_group.size() != 0):
            msg = f"`num_heads` must be divisible by `num_shards` (got `num_heads`: {self.num_heads} " \
                  f"and `num_shards`: {weights.process_group.size()}"
            logger.error(msg, ErrorCode.ATB_MODELS_EXECUTION_FAILURE)
            raise ValueError(msg)
        if config.num_key_value_heads < weights.process_group.size():
            repeat_times = weights.process_group.size() // config.num_key_value_heads
        else:
            repeat_times = 1

        self.num_heads = (self.num_heads + weights.process_group.size() - 1) // weights.process_group.size()
        if config.num_key_value_heads != config.num_attention_heads:
            self.num_key_value_heads = config.num_key_value_heads * repeat_times
            self.num_key_value_heads = self.num_key_value_heads // weights.process_group.size()
        else:
            self.num_key_value_heads = self.num_heads
        linear_names = [f'{prefix}.q_proj', f'{prefix}.k_proj', f'{prefix}.v_proj']
        pack_name = f'{prefix}.query_key_value'
        layer_prefix = '.'.join(prefix.split('.')[:-1])
        norm_name = f'{layer_prefix}.input_layernorm'
        self.pack_type = calc_linear_pack_type(weights, linear_names, norm_name, pack_name)

        if self.pack_type in [PackType.ALL_FP, PackType.ALL_W8A8, PackType.ALL_W8A8_ANTI, PackType.ALL_W8A16]:
            self.query_key_value = load_column_multi(
                config,
                prefixes=[f"{prefix}.q_proj", f"{prefix}.k_proj", f"{prefix}.v_proj"],
                weights=weights,
                head_size=self.head_size
            )
        elif self.pack_type == PackType.ALL_W8A8SC:
            pass
        else:
            pass
        self.o_proj = TensorParallelRowLinear.load(
            config,
            prefix=f"{prefix}.o_proj",
            weights=weights,
            bias=False,
            gqa_size=self.head_size,
        )
        self.num_groups = self.num_heads // self.num_key_value_heads
        self.kv_head_mapping = torch.arange(
            0, self.num_key_value_heads, dtype=torch.int32, device=weights.device
        ).repeat_interleave(self.num_groups)

        self.prefix = prefix

    def forward(
            self,
            input_args: ForwardInputArgs
    ):
        hidden_states = input_args.hidden_states
        cos = input_args.cos
        sin = input_args.sin
        cu_seqlen_prefill = input_args.cu_seqlen_prefill
        kv_cache = input_args.kv_cache
        block_tables = input_args.block_tables
        slots = input_args.slots
        input_lengths = input_args.input_lengths
        max_s = input_args.max_s
        qkv = self.query_key_value(hidden_states)
        query, kv = qkv.split(
            [
                self.head_size * self.num_heads,
                2 * self.head_size * self.num_key_value_heads,
            ],
            dim=1,
        )
        query = query.view(-1, self.num_heads, self.head_size)
        kv = kv.view(-1, 2, self.num_key_value_heads, self.head_size)

        self.rotary_emb(query, torch.select(kv, dim=1, index=0), cos, sin)

        reshape_and_cache(
            kv[:, 0], kv[:, 1], kv_cache[0], kv_cache[1], slots
        )

        # output tensor
        attn_output = torch.empty_like(query)

        # Prefill
        if cu_seqlen_prefill is not None:
            # flash attention
            flash_attn(
                query,
                torch.select(kv, dim=1, index=0),
                torch.select(kv, dim=1, index=1),
                attn_output,
                cu_seqlen_prefill,
                max_s,
                self.softmax_scale,
            )
        # Decode
        else:
            paged_attn(
                attn_output,
                query,
                kv_cache[0],
                kv_cache[1],
                self.kv_head_mapping,
                self.softmax_scale,
                block_tables,
                input_lengths,
                max_s,
            )

        return self.o_proj(attn_output.view(-1, self.num_heads * self.head_size))


class FlashDeepseekLayer(nn.Module):
    class ForwardInputArgs:
        def __init__(self,
                     hidden_states: torch.tensor,
                     residual: torch.tensor,
                     cos: torch.tensor,
                     sin: torch.tensor,
                     cu_seqlen_prefill: torch.tensor,
                     kv_cache: Tuple[torch.tensor, torch.tensor],
                     block_tables: List[torch.tensor],
                     slots: torch.tensor,
                     input_lengths: torch.tensor,
                     max_s: torch.tensor):
            self.hidden_states = hidden_states
            self.residual = residual
            self.cos = cos
            self.sin = sin
            self.cu_seqlen_prefill = cu_seqlen_prefill
            self.kv_cache = kv_cache
            self.block_tables = block_tables
            self.slots = slots
            self.input_lengths = input_lengths
            self.max_s = max_s

    def __init__(self, layer_id, config, weights):
        super().__init__()
        prefix = f"model.layers.{layer_id}"
        self.self_attn = FlashDeepseekAttention(
            prefix=f"{prefix}.self_attn", config=config, weights=weights
        )
        if (config.n_routed_experts is not None and
            layer_id >= config.first_k_dense_replace and
            layer_id % config.moe_layer_freq == 0):
            self.mlp = DeepseekMoE(prefix=f"{prefix}.mlp", config=config, weights=weights, shared_mlp_cls=DeepseekMLP)
        else:
            self.mlp = DeepseekMLP(prefix=f"{prefix}.mlp", config=config, weights=weights)
        self.input_layernorm = DeepseekRMSNorm(
            prefix=f"{prefix}.input_layernorm", weights=weights, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = DeepseekRMSNorm(
            prefix=f"{prefix}.post_attention_layernorm",
            weights=weights,
            eps=config.rms_norm_eps,
        )

    def forward(
            self,
            input_args: ForwardInputArgs
    ):
        hidden_states = input_args.hidden_states
        residual = input_args.residual
        cos = input_args.cos
        sin = input_args.sin
        cu_seqlen_prefill = input_args.cu_seqlen_prefill
        kv_cache = input_args.kv_cache
        block_tables = input_args.block_tables
        slots = input_args.slots
        input_lengths = input_args.input_lengths
        max_s = input_args.max_s
        normed_hidden_states, res = self.input_layernorm(hidden_states, residual)

        # Self Attention
        attn_output = self.self_attn(
            normed_hidden_states,
            cos,
            sin,
            cu_seqlen_prefill,
            kv_cache,
            block_tables,
            slots,
            input_lengths,
            max_s,
        )

        # faster post attention rms norm
        normed_attn_res_output, attn_res = self.post_attention_layernorm(
            attn_output, res
        )

        mlp_output = self.mlp(normed_attn_res_output)

        return mlp_output, attn_res


class FlashDeepseekModel(torch.nn.Module):
    class ForwardInputArgs:
        def __init__(self,
                    input_ids: torch.Tensor,
                    position_ids: torch.Tensor,
                    cu_seqlen_prefill: Optional[torch.Tensor],
                    kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
                    block_tables: torch.Tensor,
                    slots: torch.Tensor,
                    input_lengths: torch.Tensor,
                    max_s: int,
                    lm_head_indices: Optional[torch.Tensor] = None):
            self.input_ids = input_ids
            self.position_ids = position_ids
            self.cu_seqlen_prefill = cu_seqlen_prefill
            self.kv_cache = kv_cache
            self.block_tables = block_tables
            self.slots = slots
            self.input_lengths = input_lengths
            self.max_s = max_s
            self.lm_head_indices = lm_head_indices

    def __init__(self, config, weights):
        super().__init__()

        process_group = weights.process_group
        self.tp_rank = process_group.rank()
        self.tp_world_size = process_group.size()
        self.embed_tokens = TensorEmbedding(
            prefix="model.embed_tokens", weights=weights
        )
        self.layers = nn.ModuleList(
            [
                FlashDeepseekLayer(
                    layer_id,
                    config,
                    weights,
                )
                for layer_id in range(config.num_hidden_layers)
            ]
        )
        self.norm = DeepseekRMSNorm(
            prefix="model.norm", weights=weights, eps=config.rms_norm_eps
        )

        self.gradient_checkpointing = False

        self.head_size = self.layers[0].self_attn.head_size
        self.num_heads = self.layers[0].self_attn.num_heads
        self.num_key_value_heads = self.layers[0].self_attn.num_key_value_heads

    def forward(
            self,
            input_args: ForwardInputArgs
    ):
        input_ids = input_args.input_ids
        position_ids = input_args.position_ids
        cu_seqlen_prefill = input_args.cu_seqlen_prefill
        kv_cache = input_args.kv_cache
        block_tables = input_args.block_tables
        slots = input_args.slots
        input_lengths = input_args.input_lengths
        max_s = input_args.max_s
        hidden_states = self.embed_tokens(input_ids)

        # Get rotary cos and sin for this forward
        # Avoid to index in each layer
        cos, sin = self.layers[0].self_attn.rotary_emb.get_cos_sin(
            position_ids, max_s, hidden_states.dtype
        )

        residual = None
        for i, layer in enumerate(self.layers):
            hidden_states, residual = layer(
                hidden_states,
                residual,
                cos,
                sin,
                cu_seqlen_prefill,
                kv_cache[i],
                block_tables,
                slots,
                input_lengths,
                max_s,
            )

        hidden_states, _ = self.norm(hidden_states, residual)

        return hidden_states
    

class DeepseekEp(nn.Module):
    """
    for experts parallel.
    """

    def __init__(self, prefix, config, weights):
        super().__init__()
        expert_gate_proj = weights.get_tensor(f"{prefix}.gate_proj.weight")
        self.expert_gate_proj = nn.Parameter(expert_gate_proj)
        expert_up_proj = weights.get_tensor(f"{prefix}.up_proj.weight")
        self.expert_up_proj = nn.Parameter(expert_up_proj)
        expert_down_proj = weights.get_tensor(f"{prefix}.down_proj.weight")
        self.expert_down_proj = nn.Parameter(expert_down_proj)


class DeepseekMoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """

    def __init__(self, prefix, config, weights, shared_mlp_cls,
            gate_key="gate", shared_expert_key="shared_experts"):
        super().__init__()
        self.config = config
        self.hidden_dim = self.config.hidden_size
        self.num_experts_per_tok = config.num_experts_per_tok
        self.num_experts = config.n_routed_experts

        self.ep = weights.mapping.has_moe_ep()
        if self.ep:
            self.rank = weights.mapping.moe_ep.rank
            self.world_size = weights.mapping.moe_ep.group_size
        else:
            if weights.mapping.has_moe_tp():
                self.rank = weights.mapping.moe_tp.rank
                self.world_size = weights.mapping.moe_tp.group_size
            else:
                self.rank = weights.mapping.mlp_tp.rank
                self.world_size = weights.mapping.mlp_tp.group_size

        self.expert_lists = []
        if self.ep:
            self.expert_lists = assign(config.n_routed_experts, self.world_size)
        else:
            self.expert_lists = [[i for i in range(config.n_routed_experts)] for j in range(self.world_size)]

        self.device_expert = [i for i in range(self.config.n_routed_experts)] if not self.ep else \
            assign(self.config.n_routed_experts, weights.mapping.world_size)[weights.mapping.moe_ep.rank]
        temp_list = [j for j in range(config.n_routed_experts)]
        temp_list = temp_list[self.device_expert[0]:] + temp_list[:self.device_expert[0]]

        expert_prefix = f"{prefix}.experts"
        if hasattr(config, "topk_method") and config.topk_method == "noaux_tc":
            self.gate = FastLinear.load(
                prefix=f"{prefix}.{gate_key}", weights=weights, bias=True, bias_name="e_score_correction_bias")
            self.gate.bias.data = self.gate.bias.data[temp_list]
        else:
            self.gate = FastLinear.load(prefix=f"{prefix}.{gate_key}", weights=weights, bias=False)
        if (not hasattr(config, "ep_level")) or config.ep_level != 2:
            self.gate.weight.data = self.gate.weight.data[temp_list]

        linear_names = [f'{expert_prefix}.0.up_proj', f'{expert_prefix}.0.gate_proj']
        pack_name = f'{expert_prefix}.0.gate_up_proj'
        layer_prefix = '.'.join(prefix.split('.')[:-1])
        norm_name = f'{layer_prefix}.post_attention_layernorm'
        self.pack_type = calc_linear_pack_type(weights, linear_names, norm_name, pack_name)

        if self.ep:
            weights.switch_process_group(ProcessGroupType.MOE_EP)

        if self.pack_type in [
            PackType.ALL_FP, PackType.ALL_W8A8, PackType.ALL_W8A8_ANTI, PackType.ALL_W4A16,
            PackType.ALL_W4A16_ANTI, PackType.ALL_W8A16, PackType.ALL_W8A16_ANTI,
            PackType.MIX_W8A8_DYNAMIC, PackType.MIX_W8A8_DYNAMIC_ANTI,
            PackType.ALL_W8A8_DYNAMIC, PackType.ALL_W8A8_DYNAMIC_ANTI
        ]:

            self.gate_up_proj = nn.ModuleList()
            for i in self.expert_lists[self.rank]:
                self.gate_up_proj.append(load_column_multi(
                    config,
                    prefixes=[f"{expert_prefix}.{i}.gate_proj", f"{expert_prefix}.{i}.up_proj"],
                    weights=weights,
                    head_size=1,
                ))
        elif self.pack_type in [PackType.ALL_W8A8SC, PackType.ALL_W8A8SC_ANTI]:
            self.gate_up_proj = nn.ModuleList()
            for i in self.expert_lists[self.rank]:
                self.gate_up_proj.append(TensorParallelColumnLinear.load(
                    config,
                    prefix=f"{expert_prefix}.{i}.gate_up_proj",
                    weights=weights,
                    bias=False,
                ))
        else:
            self.gate_proj = nn.ModuleList()
            for i in self.expert_lists[self.rank]:
                self.gate_proj.append(TensorParallelColumnLinear.load(
                    config,
                    prefix=f"{expert_prefix}.{i}.gate_proj",
                    weights=weights,
                    bias=False,
                ))
            self.up_proj = nn.ModuleList()
            for i in self.expert_lists[self.rank]:
                self.up_proj.append(TensorParallelColumnLinear.load(
                    config,
                    prefix=f"{expert_prefix}.{i}.up_proj",
                    weights=weights,
                    bias=False,
                ))

        self.down_proj = nn.ModuleList()
        for i in self.expert_lists[self.rank]:
            self.down_proj.append(TensorParallelRowLinear.load(
                config,
                prefix=f"{expert_prefix}.{i}.down_proj",
                weights=weights,
                bias=False,
            ))
        self.intermediate_size = ((config.intermediate_size + self.world_size - 1) // self.world_size)

        if self.ep:
            weights.switch_process_group(ProcessGroupType.MLP)
            if hasattr(config, "ep_level") and config.ep_level == 2:
                weights.switch_process_group(ProcessGroupType.MOE_EP)

        if config.n_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            shared_expert_prefix = f"{prefix}.{shared_expert_key}"
            self.shared_experts = shared_mlp_cls(
                prefix=shared_expert_prefix,
                config=config,
                weights=weights,
                intermediate_size=intermediate_size
            )

    def forward(self, hidden_states):
        identity = hidden_states
        orig_shape = hidden_states.shape
        topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        flat_topk_idx = topk_idx.view(-1)
        y = self.moe_infer(hidden_states, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
        if self.config.n_shared_experts is not None:
            y = y + self.shared_experts(identity)
        return y

    @torch.no_grad()
    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
        token_idxs = idxs // self.num_experts_per_tok
        for i, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
            if start_idx == end_idx:
                continue
            expert = self.experts[i]
            exp_token_idx = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idx]
            expert_out = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            device = expert_cache.device
            expert_cache_cpu = expert_cache.cpu()
            expert_cache_cpu.scatter_reduce_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]).cpu(),
                                             expert_out.cpu(), reduce='sum')
            expert_cache = expert_cache_cpu.to(device=device)
        return expert_cache